"""
使用ffmpeg进行rtmp推流
"""
import os, sys

parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 4)))
sys.path.insert(0, parent_path)

import subprocess,multiprocessing
import numpy as np
from core.utils.logger import logging

class RtmpStream():
    def __init__(self, out_url, width, height, fps) -> None:
        self.url = out_url
        self.width = width
        self.height = height
        self.fps = fps
        self.command = self._build_command()
        self.pipe = subprocess.Popen(self.command,stdin=subprocess.PIPE,)

    def put(self,data):
        try:
            self.pipe.stdin.write(data.tostring())
            return True
        except Exception as e:
            logging.error("|RtmpStream put fail|")
            logging.error(e)
            return False
    
    def start(self):
        pass

    def close(self):
        self.pipe.terminate()
        # self.pipe.wait()
        # os.killpg(self.pipe.pid,signal.SIGTERM)

    def _build_command(self):
        return ['ffmpeg',
                '-y','-an',
                '-f', 'rawvideo',
                '-vcodec', 'rawvideo',
                '-pix_fmt', 'bgr24',
                '-s', "{}x{}".format(self.width, self.height),
                '-r', str(self.fps),
                '-i', '-',
                '-c:v', 'libx264',
                '-b:v', "2000k",
                '-pix_fmt', 'yuv420p',
                '-preset', 'fast',
                '-tune', 'zerolatency',
                '-f', 'flv',
                self.url]
    

class RtmpStreamMultiprocessing():
    def __init__(self, out_url, width, height, fps) -> None:
        self.url = out_url
        self.width = width
        self.height = height
        self.fps = fps
        self.command = self._build_command()
        self.buffer = multiprocessing.Queue(2)
        self.buffer.put(np.ones(shape=[height,width,3],dtype=np.uint8)*(0,0,255))
        self.status = multiprocessing.Value("i",1) # 0:standby, 1:ready, 2:running
        self.pipe = subprocess.Popen(self.command,stdin=subprocess.PIPE,)
        self.run_thread = multiprocessing.Process(target=RtmpStreamMultiprocessing.run,args=(self.pipe, self.buffer,self.status))
        

    @staticmethod
    def run(pipe, buffer:multiprocessing.Queue,status:multiprocessing.Value): # type: ignore
        assert status.value == 1, logging.error("|rtmp stream not ready|")
        with status.get_lock():
            status.value = 2 # running
        # pipe = subprocess.Popen(command,stdin=subprocess.PIPE,)
        while status.value == 2:
            try:
                data = buffer.get(timeout=20)
            except:
                logging.error("|rtmp_steam 推流失败, 无法从buffer中get数据|")
                break
            pipe.stdin.write(data.tostring())
        pipe.terminate()
        logging.info("|RtmpStream running done|")

    def put(self,data):
        try:
            self.buffer.put_nowait(data)
        except:
            pass
        if self.status.value == 2:
            return True
        else:
            return False

    def start(self,):
        self.run_thread.start()

    def get_status(self):
        return self.status.value

    def close(self):
        with self.status.get_lock():
            self.status.value = 0
        self.pipe.terminate()

    def _build_command(self):
        return ['ffmpeg',
                '-y','-an',
                '-f', 'rawvideo',
                '-vcodec', 'rawvideo',
                '-pix_fmt', 'bgr24',
                '-s', "{}x{}".format(self.width, self.height),
                '-r', str(self.fps),
                '-i', '-',
                '-c:v', 'libx264',
                '-pix_fmt', 'yuv420p',
                '-preset', 'fast',
                '-tune', 'zerolatency',
                '-f', 'flv',
                self.url]
    

if __name__ == "__main__":
    import cv2, time, datetime
    from utils.video_reader import VideoReader
    video_file_path = "rtmp://218.94.155.86:1936/live/livestream_74_1"
    reader = VideoReader(mode="block")
    reader.init_cap(video_file_path)
    reader.start()

    rtmp_stream = RtmpStream(out_url="rtmp://218.94.155.86:1936/live/livestream_74_3",
                            width=reader.width,
                            height=reader.height,
                            fps=25)
    
    wait_time = 1/rtmp_stream.fps
    start_time = time.time()

    while 1: 
        runing_start = time.time()
        ret, frame = reader.read()
        # cv2.putText(frame,time.strftime('%Y-%m-%d_%H-%M-%S',time.localtime()),[50,50],cv2.FONT_ITALIC,1,(0,255,255)) 
        rtmp_stream.put(frame)
        # time.sleep(0.01)
        # time.sleep(max(wait_time-(time.time()-runing_start),0))
        # ffmpeg -y -an -f rawvideo -vcodec rawvideo -pix_fmt bgr24 -s 1920x1080 -r 25 -i - -c:v libx264 -b:v 2000k -pix_fmt yuv420p -preset fast -tune zerolatency -f flv rtmp://218.94.155.86:1936/live/livestream_74_3